/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <melon/synchronized.h>

/* `SynchronizedPtr` is a variation on the `Synchronized` idea that's useful for
 * some cases where you want to protect a pointed-to object (or an object within
 * some pointer-like wrapper). If you would otherwise need to use
 * `Synchronized<smart_ptr<Synchronized<T>>>` consider using
 * `SynchronizedPtr<smart_ptr<T>>`as it is a bit easier to use and it works when
 * you want the `T` object at runtime to actually a subclass of `T`.
 *
 * You can access the contained `T` with `.rlock()`, and `.wlock()`, and the
 * pointer or pointer-like wrapper with `.wlockPointer()`. The corresponding
 * `with...` methods take a callback, invoke it with a `T const&`, `T&` or
 * `smart_ptr<T>&` respectively, and return the callback's result.
 */
namespace melon {
    template<typename LockHolder, typename Element>
    struct SynchronizedPtrLockedElement {
        explicit SynchronizedPtrLockedElement(LockHolder &&holder)
            : holder_(std::move(holder)) {
        }

        Element &operator*() const { return **holder_; }

        Element *operator->() const { return &**holder_; }

        explicit operator bool() const { return static_cast<bool>(*holder_); }

    private:
        LockHolder holder_;
    };

    template<typename PointerType, typename MutexType = SharedMutex>
    class SynchronizedPtr {
        using inner_type = Synchronized<PointerType, MutexType>;
        inner_type inner_;

    public:
        using pointer_type = PointerType;
        using element_type = typename std::pointer_traits<pointer_type>::element_type;
        using const_element_type = typename std::add_const<element_type>::type;
        using read_locked_element = SynchronizedPtrLockedElement<
            typename inner_type::ConstLockedPtr,
            const_element_type>;
        using write_locked_element = SynchronizedPtrLockedElement<
            typename inner_type::LockedPtr,
            element_type>;
        using write_locked_pointer = typename inner_type::LockedPtr;

        template<typename... Args>
        explicit SynchronizedPtr(Args... args)
            : inner_(std::forward<Args>(args)...) {
        }

        SynchronizedPtr() = default;

        SynchronizedPtr(SynchronizedPtr const &) = default;

        SynchronizedPtr(SynchronizedPtr &&) = default;

        SynchronizedPtr &operator=(SynchronizedPtr const &) = default;

        SynchronizedPtr &operator=(SynchronizedPtr &&) = default;

        // Methods to provide appropriately locked and const-qualified access to the
        // element.

        read_locked_element rlock() const {
            return read_locked_element(inner_.rlock());
        }

        template<class Function>
        auto withRLock(Function &&function) const {
            return function(*rlock());
        }

        write_locked_element wlock() { return write_locked_element(inner_.wlock()); }

        template<class Function>
        auto withWLock(Function &&function) {
            return function(*wlock());
        }

        // Methods to provide write-locked access to the pointer. We deliberately make
        // it difficult to get a read-locked pointer because that provides read-locked
        // non-const access to the element, and the purpose of this class is to
        // discourage that.
        write_locked_pointer wlockPointer() { return inner_.wlock(); }

        template<class Function>
        auto withWLockPointer(Function &&function) {
            return function(*wlockPointer());
        }
    };
} // namespace melon
